#include "StdAfx.h"
#include "D3D9ShaderCompiler.h"

#include "StringUtility.h"


/**
* The D3DX/HLSL shader compiler.
*/ 
using namespace Common;
using namespace RenderSystem;

/**
* An implementation of the D3DX include interface to access a ShaderCompilerEnvironment.
*/
class D3DIncludeEnvironment : public ID3DXInclude
{
public:

	//constructor
	D3DIncludeEnvironment( ShaderCompilerEnvironment& rInEnvironment, ShaderCompiler& rShaderCompiler)
		:m_rEnvironment(rInEnvironment), m_ShaderCompiler(rShaderCompiler)
	{
	}

    STDMETHOD(Open)( D3DXINCLUDE_TYPE IncludeType, LPCSTR pFileName, LPCVOID pParentData, LPCVOID *ppData, UINT *pBytes )
	{
		TCHAR szFileName[MAX_PATH];
		StringUtility::MByteToTCHAR( pFileName, szFileName, MAX_PATH );

		String strFileName( szFileName );
		String strFileContents;

		ShaderCompilerEnvironment::CompileMapIterator itrIncludeFiles = m_rEnvironment.IncludeFiles.find( strFileName );
		if( itrIncludeFiles != m_rEnvironment.IncludeFiles.end() )
		{
			strFileContents = itrIncludeFiles->second;
		}
		else
		{
			strFileContents = m_ShaderCompiler.LoadShaderSourceFile( strFileName.c_str() );
		}

		// Convert the file contents to ANSI and new a buffer

		char* pszFileContents = new char[ strFileContents.length() + 1 ];

		// Write the result to the output parameters.
		*ppData = (LPCVOID)pszFileContents;
		*pBytes = strFileContents.size();

		return D3D_OK;
	}

	STDMETHOD(Close)(LPCVOID Data)
	{
		delete [] Data;
		return D3D_OK;
	}

private:

	ShaderCompilerEnvironment&				m_rEnvironment;
	ShaderCompiler&							m_ShaderCompiler;
};

bool D3D9ShaderCompiler::_CompileShader(
	const TCHAR* pszSourceFilename,
	const TCHAR* pszFunctionName,
	const ShaderCompilerEnvironment& rEnvironment,
	const ShaderProfile& rShaderProfile, 
	ShaderCompilerOutput& rOutput
	)
{
	//handle definitions
	int iDefCount = rEnvironment.Definitions.size();
	D3DXMACRO* pShaderMacros = new D3DXMACRO[ iDefCount + 1 ];

	DWORD dwCompileFlags = D3DXSHADER_ENABLE_BACKWARDS_COMPATIBILITY;

#ifdef _DEBUG
		//add the debug flags
		dwCompileFlags |= D3DXSHADER_DEBUG | D3DXSHADER_SKIPOPTIMIZATION;
#endif

	ShaderCompilerEnvironment::CompileMapConstIterator itrDef = rEnvironment.Definitions.begin();
	for( int iIdx = 0; itrDef != rEnvironment.Definitions.end(); ++itrDef, ++iIdx )
	{
		D3DXMACRO* macro = &pShaderMacros[iIdx];
		macro->Name = itrDef->first.c_str();
		macro->Definition = itrDef->second.c_str();
	}
	pShaderMacros[ iDefCount ].Definition = NULL;
	pShaderMacros[ iDefCount ].Name	= NULL;

	ShaderFileCachesIterator itrShaderFile = m_ShaderFileCaches.find( pszSourceFilename );
	
	String strShaderContent;
	if( itrShaderFile == m_ShaderFileCaches.end() )
		strShaderContent = this->LoadShaderSourceFile( pszSourceFilename );
	else
		strShaderContent = itrShaderFile->second;

	D3DIncludeEnvironment d3dIncludes( const_cast<ShaderCompilerEnvironment&>( rEnvironment ), *this );

	LPD3DXBUFFER pShader;
	LPD3DXBUFFER pErrorMsgs;
	LPD3DXCONSTANTTABLE pConstantTable;

	if( D3DXCompileShader( strShaderContent.c_str(), 
		strShaderContent.size(), 
		pShaderMacros,
		&d3dIncludes,
		pszFunctionName,
		rShaderProfile.GetProfile().c_str(),
		dwCompileFlags,
		&pShader,
		&pErrorMsgs,
		&pConstantTable ) != D3D_OK )
	{
		//report error info
	}

	//TODO:generate a compile information into compile queue.
	//here just compile it directly
	UINT uNumShaderBytes = pShader->GetBufferSize();
	rOutput.m_strCodes.append( (TCHAR*)pShader->GetBufferPointer() );

	// Read the constant table output from the shader bytecode.
	D3DXCONSTANTTABLE_DESC objConstantTableDesc;
	pConstantTable->GetDesc(&objConstantTableDesc);

	// Read the constant descriptions out of the shader bytecode.
	for(UINT uConstantIndex = 0; uConstantIndex < objConstantTableDesc.Constants; uConstantIndex++)
	{
		// Read the constant description.
		D3DXHANDLE hConstantHandle = pConstantTable->GetConstant(NULL, uConstantIndex);
		D3DXCONSTANT_DESC objConstantDesc;
		UINT uNumConstants = 1;
		pConstantTable->GetConstantDesc(hConstantHandle, &objConstantDesc, &uNumConstants);

		if(objConstantDesc.RegisterSet == D3DXRS_SAMPLER)
		{
			rOutput.m_ParameterMap.AddParameterAllocation(
				objConstantDesc.Name,
				0,
				objConstantDesc.RegisterIndex,
				objConstantDesc.RegisterCount,
				objConstantDesc.RegisterIndex
				);
		}
		else
		{
			rOutput.m_ParameterMap.AddParameterAllocation(
				objConstantDesc.Name,
				0,
				objConstantDesc.RegisterIndex * sizeof(FLOAT) * 4,
				objConstantDesc.RegisterCount * sizeof(FLOAT) * 4,
				0
				);
		}
	}

	// Disassemble the shader to determine the instruction count.

	return true;
}

/*
bool D3D9ShaderCompiler::CompileShader(
	const TCHAR* pszSourceFilename,
	const TCHAR* pszFunctionName,
	const ShaderCompilerEnvironment& rEnvironment,
	ShaderCompilerOutput& rOutput
	)
{
	TArray<D3D9ConstantDesc> Constants;
	FString DisassemblyString;

	// Translate the input environment's definitions to D3DXMACROs.
	TArray<D3DXMACRO> Macros;
	for(TMap<FName,FString>::TConstIterator DefinitionIt(Environment.Definitions);DefinitionIt;++DefinitionIt)
	{
		FString Name = DefinitionIt.Key().ToString();
		FString Definition = DefinitionIt.Value();

		D3DXMACRO* Macro = new(Macros) D3DXMACRO;
		Macro->Name = new ANSICHAR[Name.Len() + 1];
		appStrncpyANSI( (ANSICHAR*)Macro->Name, TCHAR_TO_ANSI(*Name), Name.Len()+1 );
		Macro->Definition = new ANSICHAR[Definition.Len() + 1];
		appStrncpyANSI( (ANSICHAR*)Macro->Definition, TCHAR_TO_ANSI(*Definition), Definition.Len()+1 );
	}

	// set the COMPILER type
	D3DXMACRO* Macro = new(Macros) D3DXMACRO;
#define COMPILER_NAME "COMPILER_HLSL"
	Macro->Name = appStrcpyANSI(
		new ANSICHAR[strlen(COMPILER_NAME) + 1], 
		strlen(COMPILER_NAME) + 1,
		COMPILER_NAME
		);
	Macro->Definition = appStrcpyANSI(new ANSICHAR[2], 2, "1");

	DWORD CompileFlags = D3DXSHADER_ENABLE_BACKWARDS_COMPATIBILITY;
	if (DEBUG_SHADERS) 
	{
		//add the debug flags
		CompileFlags |= D3DXSHADER_DEBUG | D3DXSHADER_SKIPOPTIMIZATION;
	}

	for(INT FlagIndex = 0;FlagIndex < Environment.CompilerFlags.Num();FlagIndex++)
	{
		//accumulate flags set by the shader
		CompileFlags |= TranslateCompilerFlag(Environment.CompilerFlags(FlagIndex));
	}

	bool bShaderCompileFailed = TRUE;
	const FString ShaderPath = appShaderDir();
	const FString PreprocessorOutputDir = ShaderPath * ShaderPlatformToText((EShaderPlatform)Target.Platform) * FString(ShaderSubDir);

	TCHAR ShaderProfile[32];
	TArray<FString> FilteredErrors;
	{
		//set defines and profiles for the appropriate shader paths
		if (Target.Platform == SP_PCD3D_SM3)
		{
			if (Target.Frequency == SF_Pixel)
			{
				appStrcpy(ShaderProfile, TEXT("ps_3_0"));
			}
			else
			{
				appStrcpy(ShaderProfile, TEXT("vs_3_0"));
			}

			// Set SM3_PROFILE, which indicates that we are compiling for sm3.  
			D3DXMACRO* MacroSM3Profile = new(Macros) D3DXMACRO;
			MacroSM3Profile->Name = appStrcpyANSI(
				new ANSICHAR[strlen("SM3_PROFILE") + 1], 
				strlen("SM3_PROFILE") + 1,
				"SM3_PROFILE"
				);
			MacroSM3Profile->Definition = appStrcpyANSI(new ANSICHAR[2], 2, "1");
		}
		else if (Target.Platform == SP_PCD3D_SM2)
		{
			if (Target.Frequency == SF_Pixel)
			{
				appStrcpy(ShaderProfile, TEXT("ps_2_0"));
			}
			else
			{
				appStrcpy(ShaderProfile, TEXT("vs_2_0"));
			}

			// Set the SUPPORTS_FP_BLENDING shader define.
			// It defaults to TRUE in Definitions.usf, so it's not necessary to set it in that case.
			D3DXMACRO* MacroManualFPBlend = new(Macros) D3DXMACRO;
			MacroManualFPBlend->Name = appStrcpyANSI(
				new ANSICHAR[strlen("SUPPORTS_FP_BLENDING") + 1], 
				strlen("SUPPORTS_FP_BLENDING") + 1,
				"SUPPORTS_FP_BLENDING"
				);
			MacroManualFPBlend->Definition = appStrcpyANSI(new ANSICHAR[2], 2, "0");

			// Set SM2_PROFILE, which indicates that we are compiling for sm2.  
			// This is needed to work around some HLSL compiler bugs with this profile.
			D3DXMACRO* MacroSM2Profile = new(Macros) D3DXMACRO;
			MacroSM2Profile->Name = appStrcpyANSI(
				new ANSICHAR[strlen("SM2_PROFILE") + 1], 
				strlen("SM2_PROFILE") + 1,
				"SM2_PROFILE"
				);
			MacroSM2Profile->Definition = appStrcpyANSI(new ANSICHAR[2], 2, "1");
		}

		// Add a define that indicates whether the compiler supports attributes such as [unroll]
		INT AttributeMacroIndex = Macros.Num();
		D3DXMACRO* MacroCompilerSupportsAttributes = new(Macros) D3DXMACRO;
		MacroCompilerSupportsAttributes->Name = appStrcpyANSI(
			new ANSICHAR[strlen("COMPILER_SUPPORTS_ATTRIBUTES") + 1], 
			strlen("COMPILER_SUPPORTS_ATTRIBUTES") + 1,
			"COMPILER_SUPPORTS_ATTRIBUTES"
			);
		MacroCompilerSupportsAttributes->Definition = appStrcpyANSI(new ANSICHAR[2], 2, "1");

		// Terminate the Macros list.
		D3DXMACRO* TerminatorMacro = new(Macros) D3DXMACRO;
		TerminatorMacro->Name = NULL;
		TerminatorMacro->Definition = NULL;

		// If this is a multi threaded compile, we must use the worker, otherwise compile the shader directly.
		if (GShaderCompilingThreadManager->IsMultiThreadedCompile())
		{
			bShaderCompileFailed = !D3D9CompileShaderThroughWorker(
				ThreadId,
				SourceFilename,
				TCHAR_TO_ANSI(FunctionName),
				TCHAR_TO_ANSI(ShaderProfile),
				TCHAR_TO_ANSI(*ShaderPath),
				CompileFlags,
				Environment,
				Output,
				Macros,
				Constants,
				DisassemblyString,
				FilteredErrors
				);
		}
		else
		{
			bShaderCompileFailed = !D3D9CompileShaderThroughDll(
				SourceFilename,
				FunctionName,
				ShaderProfile,
				CompileFlags,
				Environment,
				Output,
				Macros,
				AttributeMacroIndex,
				Constants,
				DisassemblyString,
				FilteredErrors
				);
		}
	}

	// If we are dumping out preprocessor data
	// @todo - also dump out shader data when compilation fails
	if (bDebugDump)
	{
		// just in case the preprocessed shader dir has not been created yet
		GFileManager->MakeDirectory( *PreprocessorOutputDir, TRUE );

		// save out include files from the environment definitions
		// Note: Material.usf and VertexFactory.usf vary between materials/vertex factories
		// this is handled because fxc will search for the includes in the same directory as the main shader before searching the include path 
		// otherwise it would find a stale Material.usf and VertexFactory.usf left behind by other platforms
		for(TMap<FString,FString>::TConstIterator IncludeIt(Environment.IncludeFiles); IncludeIt; ++IncludeIt)
		{
			FString IncludePath = PreprocessorOutputDir * IncludeIt.Key();
			appSaveStringToFile(IncludeIt.Value(), *IncludePath);
		}

		const FString SourceFile = LoadShaderSourceFile(SourceFilename);
		const FString SaveFileName = PreprocessorOutputDir * SourceFilename;
		appSaveStringToFile(SourceFile, *(SaveFileName + TEXT(".usf")));

		// dump the preprocessed shader
		FD3DIncludeEnvironment IncludeEnvironment(Environment);
		D3D9PreProcessShader(SourceFilename, SourceFile, Macros, IncludeEnvironment, *PreprocessorOutputDir);

		FString AbsoluteShaderPath = FString(appBaseDir()) * PreprocessorOutputDir;
		FString AbsoluteIncludePath = FString(appBaseDir()) * ShaderPath;
		// get the fxc command line
		FString FXCCommandline = D3D9CreateShaderCompileCommandLine(
			AbsoluteShaderPath * SourceFilename + TEXT(".usf"), 
			AbsoluteIncludePath, 
			FunctionName, 
			ShaderProfile,
			&Macros(0),
			CompileFlags);

		appSaveStringToFile(FXCCommandline, *(SaveFileName + TEXT(".bat")));
	}

	// Free temporary strings allocated for the macros.
	for(INT MacroIndex = 0;MacroIndex < Macros.Num();MacroIndex++)
	{
		delete Macros(MacroIndex).Name;
		delete Macros(MacroIndex).Definition;
	}

	for (INT ErrorIndex = 0; ErrorIndex < FilteredErrors.Num(); ErrorIndex++)
	{
		const FString& CurrentError = FilteredErrors(ErrorIndex);
		FShaderCompilerError NewError;
		// Extract the filename and line number from the shader compiler error message for PC whose format is:
		// "d:\UnrealEngine3\Binaries\memory(210,5): error X3004: undeclared identifier 'LightTransfer'"
		INT FirstParenIndex = CurrentError.InStr(TEXT("("));
		INT LastParenIndex = CurrentError.InStr(TEXT("):"));
		if (FirstParenIndex != INDEX_NONE 
			&& LastParenIndex != INDEX_NONE
			&& LastParenIndex > FirstParenIndex)
		{
			NewError.ErrorFile = CurrentError.Left(FirstParenIndex);
			NewError.ErrorLineString = CurrentError.Mid(FirstParenIndex + 1, LastParenIndex - FirstParenIndex - appStrlen(TEXT("(")));
			NewError.StrippedErrorMessage = CurrentError.Right(CurrentError.Len() - LastParenIndex - appStrlen(TEXT("):")));
		}
		else
		{
			NewError.StrippedErrorMessage = CurrentError;
		}
		Output.Errors.AddItem(NewError);
	}

	UINT NumInstructions = 0;
	if (!bShaderCompileFailed)
	{
		// Extract the instruction count from the disassembly text.
		if (Parse(
			*DisassemblyString,
			TEXT("// approximately "),
			(DWORD&)NumInstructions
			))
		{
			// Fail compilation if the shader used more than the minimum number of instructions available in Shader Model 3.0 for both pixel and vertex shaders.
			// Instruction count for SM 3.0 is only validated at runtime (CreateVertex/PixelShader), this check catches the incompatibility early,
			// Since some older SM3.0 cards only support 544 instructions (Geforce 7900 GT, others) and we only compile one version for all SM3 targets.
			if (NumInstructions > 512 && Target.Platform == SP_PCD3D_SM3)
			{
				bShaderCompileFailed = TRUE;
				FShaderCompilerError NewError;
				NewError.StrippedErrorMessage = FString::Printf(TEXT("Shader uses too many instructions for Shader Model 3.0! (Used %u, Minimum guaranteed in SM3.0 is 512)"), NumInstructions);
				Output.Errors.AddItem(NewError);
			}
		}
	}

	// if the shader failed to compile we need to tell the caller that Compiling the shader failed
	if( bShaderCompileFailed == TRUE )
	{
		return FALSE;
	}

	// If the shader compile succeeded then this should contain the bytecode
	check(Output.Code.Num() > 0);

	// Map each constant in the table.
	for(INT ConstantIndex = 0;ConstantIndex < Constants.Num();ConstantIndex++)
	{
		// Read the constant description.
		const FD3D9ConstantDesc& Constant = Constants(ConstantIndex);

		if(Constant.bIsSampler)
		{
			Output.ParameterMap.AddParameterAllocation(
				ANSI_TO_TCHAR(Constant.Name),
				0,
				Constant.RegisterIndex,
				Constant.RegisterCount,
				Constant.RegisterIndex
				);
		}
		else
		{
			Output.ParameterMap.AddParameterAllocation(
				ANSI_TO_TCHAR(Constant.Name),
				0,
				Constant.RegisterIndex * sizeof(FLOAT) * 4,
				Constant.RegisterCount * sizeof(FLOAT) * 4,
				0
				);
		}
	}

	// Pass the target through to the output.
	Output.Target = Target;
	Output.NumInstructions = NumInstructions;

	if(bDebugDump)
	{
		appSaveStringToFile( DisassemblyString, *(PreprocessorOutputDir * SourceFilename + TEXT(".asm")) );
	}

	return TRUE;
}
*/
